import random
import torch
# rom d2l import torch as d2l

if torch.cuda.is_available():
    print('Current device:', torch.cuda.current_device())
    print('Device name:', torch.cuda.get_device_name(0))
    # 设置多个可用GPU时，将索引从0开始设置
else:
    print('Using CPU!')

# 如果需要指定显卡：
torch.cuda.set_device(0)  # 替换为要使用的GPU索引